# #!/usr/bin/env python3
# import rospy
# import matplotlib.pyplot as plt
# import matplotlib.animation as animation
# from std_msgs.msg import Float32MultiArray

# # 存储 pitch 数据
# pitch_data = []

# def callback(data):
#     global pitch_data
#     if len(data.data) >= 2:
#         pitch = data.data[1]  # 话题数据格式为 [roll, pitch]
#         pitch_data.append(-pitch)
#         if len(pitch_data) > 100:  # 只保留最近 100 个数据点
#             pitch_data.pop(0)

# def animate(i):
#     plt.cla()
#     plt.plot(pitch_data, label="Pitch Angle (°)")
#     plt.xlabel("Time (frames)")
#     plt.ylabel("Pitch Angle (°)")
#     plt.title("Real-time Pitch Angle Visualization")
#     plt.legend()
#     plt.grid()

# def listener():
#     rospy.init_node('pitch_visualizer', anonymous=True)
#     rospy.Subscriber("/robot/tilt_angles", Float32MultiArray, callback)
    
#     fig = plt.figure()
#     ani = animation.FuncAnimation(fig, animate, interval=100)  # 100ms 更新一次
#     plt.show()

# if __name__ == '__main__':
#     try:
#         listener()
#     except rospy.ROSInterruptException:
#         pass


#!/usr/bin/env python3
import rospy
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from std_msgs.msg import Float32MultiArray
import numpy as np

# 存储原始 pitch 数据
pitch_data = []
filtered_pitch_data = []

# 滑动平均滤波窗口大小
WINDOW_SIZE = 5
# 低通滤波参数 (α 值越小，平滑效果越明显)
ALPHA = 0.1

# 用于存储 EMA 计算的值
ema_pitch = None  

def moving_average_filter(data_list, window_size):
    """计算滑动平均值"""
    if len(data_list) < window_size:
        return np.mean(data_list)  # 如果数据点少于窗口大小，直接取均值
    return np.mean(data_list[-window_size:])  # 取最近 window_size 个数据的均值

def low_pass_filter(new_value, last_value, alpha):
    """一阶低通滤波"""
    return alpha * new_value + (1 - alpha) * last_value

def callback(data):
    global pitch_data, filtered_pitch_data, ema_pitch

    if len(data.data) >= 2:
        pitch = data.data[1]  # 话题数据格式为 [roll, pitch]
        pitch_data.append(pitch)

        # 滑动平均滤波
        smoothed_pitch = moving_average_filter(pitch_data, WINDOW_SIZE)

        # 一阶低通滤波
        if ema_pitch is None:  # 初始化
            ema_pitch = pitch
        else:
            ema_pitch = low_pass_filter(pitch, ema_pitch, ALPHA)

        # 存储平滑后的数据
        filtered_pitch_data.append(ema_pitch)

        # 限制数据长度
        if len(pitch_data) > 100:
            pitch_data.pop(0)
        if len(filtered_pitch_data) > 100:
            filtered_pitch_data.pop(0)

def animate(i):
    plt.cla()
    plt.plot(pitch_data, label="Raw Pitch (°)", color="red", alpha=0.5)  # 原始数据
    plt.plot(filtered_pitch_data, label="Filtered Pitch (°)", color="blue")  # 平滑数据
    plt.xlabel("Time (frames)")
    plt.ylabel("Pitch Angle (°)")
    plt.title("Real-time Pitch Angle Visualization (Filtered)")
    plt.legend()
    plt.grid()

def listener():
    rospy.init_node('pitch_visualizer', anonymous=True)
    rospy.Subscriber("/robot/tilt_angles", Float32MultiArray, callback)
    
    fig = plt.figure()
    ani = animation.FuncAnimation(fig, animate, interval=100)  # 100ms 更新一次
    plt.show()

if __name__ == '__main__':
    try:
        listener()
    except rospy.ROSInterruptException:
        pass